# 该脚本用于创建ssh连接, 并通过ssh连接执行linux指令
# 创建sftp链接, 用于sftp上传下载文件
import os
import sys
import logging
from time import sleep
from paramiko import AuthenticationException, SFTPClient, Transport
from paramiko.client import SSHClient, AutoAddPolicy
from paramiko.ssh_exception import NoValidConnectionsError


class Client:
    #   address(str): 目标主机ip地址
    def __init__(self, address, port=22):
        self.ssh_invoke = None
        self.address = address
        self.port = int(port)
        self.ssh_client = SSHClient()
        self.sftp_client = None
        self.status_code = 0

    # 功能: 用于检查对象是否初始化
    # 参数: obj(object): 待检测的对象
    @classmethod
    def check_blob(cls, obj):
        if obj:
            return True

        print("Please login first.")
        return False

    # 功能: 用于登录主机, 当账号
    # 参数:
    #   username(str): 登录账号
    #   password(str): 登录密码
    #   root_pass(str): root密码
    # 返回: 不返回, 创建ssh_invoke对象
    def login(self, username, password, root_pass=""):
        try:
            # 设置允许连接known_hosts文件中的主机(默认连接不在known_hosts文件中的主机会拒绝连接抛出SSHException)
            self.ssh_client.set_missing_host_key_policy(AutoAddPolicy())
            self.ssh_client.connect(self.address, port=self.port, username=username, password=password)
            self.ssh_invoke = self.ssh_client.invoke_shell()
            self.ssh_invoke.recv(9999)
            if username == "root":
                self.execute("\r\n")
            else:
                self.execute([{"command": "su root", "exp": ":"}, root_pass or password])

            self.status_code = 1000

            handle = Transport((self.address, self.port))
            handle.connect(username=username, password=password)
            self.sftp_client = SFTPClient.from_transport(handle)
        except AuthenticationException:
            logging.warning('username or password error')

            self.status_code = 1001
        except NoValidConnectionsError:
            logging.warning('connect time out')

            self.status_code = 1002
        except Exception as e:
            logging.warning('unknow error:', e)
            print("Unexpected error:", sys.exc_info()[0])

            self.status_code = 1003

        return self.status_code

    # 功能: 此函数用于执行command参数中的命令并打印命令执行结果
    # 参数:
    #   cmd(str | list | dict): 待执行的命令或者命令集
    #   timeout(int): 执行命令后的老化时间
    # 返回(str): 执行命令的结果, 包括执行的命令
    def execute(self, cmd, exp=["$", "#"], timeout=3, verbose=False):
        if not self.check_blob(self.ssh_invoke):
            return

        if isinstance(cmd, list):
            ret = ""
            for item in cmd:
                ret += self.execute(item, exp, timeout)

            return ret
        elif isinstance(cmd, dict):
            return self.execute(cmd["command"], cmd["exp"], timeout)
        else:
            if verbose:
                print(f"cmd: {cmd}, exp: {exp}, timeout: {timeout}, resp: ", end="")

            self.ssh_invoke.send(f"{cmd}\n")
            resp = b""
            while timeout > 0:
                resp += self.ssh_invoke.recv(9999)
                exp = [exp] if isinstance(exp, str) else exp
                for item in exp:
                    if item.encode("utf-8") in resp:
                        timeout = 0
                        break

                timeout -= 1
                time.sleep(1)

            # 处理打印乱码
            for item in set(re.findall(b"(\x1b\[\?\d+.)", resp)):
                resp = resp.replace(item, b"")

            resp = resp.decode("utf-8")
            print(f'{resp}')

            return resp

    def download(self, file, target):
        if not self.sftp_client:
            raise OSError("未登录")

        return self.sftp_client.get(file, target)

    def upload(self, file, target):
        if not self.sftp_client:
            raise OSError("未登录")

        return self.sftp_client.put(file, target)

    # 功能: 用于退出ssh登录
    def logout(self):
        logging.warning('Connection is closing.')
        self.status_code = 0
        self.ssh_client.close()

    def close(self):
        self.logout()


class SftpClient:
    def __init__(self):
        self.handle = None
        self.sftp = None

    # 功能: 用于sftp登录
    # 参数:
    #   host(str): 目标主机ip地址
    #   username(str): 登录账号
    #   password(str): 登录密码
    #   port(int): 端口号
    # 说明: 这里暂时未找到切换到root的方法, 若下载文件, 需要通过命令行设置文件权限
    #   经过对比, 发现使用MorbaXterm软件, 即使ssh账号切换到root, sftp权限也并没有切换
    def sftp_login(self, host, username, password, port=22):
        self.handle = Transport((host, int(port)))
        self.handle.connect(username=username, password=password)
        self.sftp = SFTPClient.from_transport(self.handle)

    # 功能: 上传目录
    # 参数:
    #   local_dir(str): 本地文件
    #   remote_dir(str): 目标文件
    # 返回: status(bool): 状态码 result(str): 执行结果
    # 说明: 当存在一个未执行的
    def upload_files(self, local_dir, remote_dir):
        status, result = True, ""
        if SshClient.check_blob(self.sftp):
            for now_dir, _, files in os.walk(local_dir):
                for file in files:
                    status, result = self.upload_file(os.path.join(now_dir, file), os.path.join(remote_dir, file))
                    if not status:
                        return status, result
        else:
            status = False

        return status, result

    # 功能: 上传文件
    # 参数:
    #   local_dir(str): 本地文件
    #   remote_dir(str): 目标文件
    # 返回: status(bool): 状态码 result(str): 执行结果
    def upload_file(self, src, dst):
        status, result = True, ""
        if SshClient.check_blob(self.sftp) and os.path.isfile(src):
            try:
                print('正在上传:', src)
                result = self.sftp.put(src, dst.replace('\\', '/'))
                print('上传完毕:', src)
            except Exception as e:
                print(e)
                status = False
                result = str(e)
        else:
            status = False

        return status, result

    # 功能: 下载文件
    # 参数:
    #   local_dir(str): 本地文件
    #   remote_dir(str): 目标文件
    # 返回: status(bool): 状态码 result(str): 执行结果
    def download_file(self, src, dst):
        status, result = True, ""
        if SshClient.check_blob(self.sftp):
            try:
                print('正在下载:', src)
                result = self.sftp.get(src.replace('\\', '/'), dst)
                print('上传完毕:', src)
            except Exception as e:
                print(e)
                status = False
                result = str(e)
        else:
            status = False

        return status, result

    # 功能: 用于sftp登出
    def sftp_close(self):
        try:
            self.handle.close()
            self.sftp = None
        except Exception as e:
            print(e)


if __name__ == "__main__":
    client = SshClient()
    client.ssh_login("10.118.185.196", "xxx", "xxx")
    commands = ['su - root', "xxx", 'sleep 2']
    for command in commands:
        client.execute_some_command(command)


